#include "header.h"

static fftw_complex *a,*hu[1],*hk[2], *lk, *ddk[2], *Dyk, *Rk, *Dk, *D, inner, *kD;
static fftw_complex *MiRk;
static fftw_plan plan_f1, plan_b1;
static fftw_plan plan_fr1, plan_br1;
static double *k, *u, overN, q, *mk, dtau, gama, sp, num, den, err, err2, l2, tol, hmax, ml0, ml, htt;
static double coeff, aa, bb, NV, *ra[2], *rt1, *rt2;
static int N, n2, fix_flag; 
static work_ptr tmp;

void err_msg(char* str)
{
  printf("%s\n",str);
  exit(1);
}

double signum(double x) {
  if (x > 0) return  1.;
  if (x < 0) return -1.;
  return 0.;
}

double signum2(double x) {
  if (x > 0) return  1.;
  else return -1.;
}

void init(int *argc, char **argv, params_ptr in, work_ptr wrk)
{
  char str[480], line[480], value[480], param[480];
  sprintf(str, "Usage:\n\t%s input", argv[0]);
  if (*argc != 2) err_msg(str);

  sprintf(str, "%s.cfg", argv[1]);
  FILE *fh = fopen(str,"r"); 
  FILE *fhlog = fopen("run.log","w"); 
  if (fh == NULL) err_msg("Cannot open file");
  while (fgets(line, 480, fh)!=NULL) {
    sscanf(line, "%s\t%s", param, value);
    if (strcmp(param,"#runname=") == 0) sprintf(in->runname,"%s", value);
    if (strcmp(param,"#resname=") == 0) sprintf(in->resname,"%s", value);
    if (strcmp(param,"#npoints=") == 0) in->n = atoi(value);
    if (strcmp(param,"#gravity=") == 0) in->g = atof(value);
    if (strcmp(param,"#velocit=") == 0) in->c = atof(value);
    if (strcmp(param,"#delta_t=") == 0) in->dtau = atof(value);
    if (strcmp(param,"#toleran=") == 0) in->tol = atof(value);
    if (strcmp(param,"#perturb=") == 0) in->alfa = atof(value);
    if (strcmp(param,"#io_time=") == 0) in->t_skip = atoi(value);
    if (strcmp(param,"#hmax_tt=") == 0) htt = atof(value);
    if (strcmp(param,"#find_xu=") == 0) in->fix_flag = atoi(value);
  }
  fprintf(fhlog, "runname = %s\n", in->runname );
  fprintf(fhlog, "resname = %s\n", in->resname );
  fprintf(fhlog, "npoints = %d\n", in->n );
  fprintf(fhlog, "gravity = %f\n", in->g );
  fprintf(fhlog, "velocit = %f\n", in->c );
  fprintf(fhlog, "toleran = %f\n", in->tol );
  fprintf(fhlog, "delta_t = %f\n", in->dtau );
  fprintf(fhlog, "perturb = %f\n", in->alfa );
  fprintf(fhlog, "hmax_tt = %f\n", htt );
  fprintf(fhlog, "io_time = %d\n", in->t_skip );
  fclose(fh); fclose(fhlog);
  N = in->n; n2 = N/2; fix_flag = in->fix_flag;
  wrk->Y   = fftw_malloc(sizeof(fftw_complex)*(in->n));
  wrk->Yk   = fftw_malloc(sizeof(fftw_complex)*(in->n));

  a  = fftw_malloc(sizeof(fftw_complex)*(in->n));
  if (fix_flag != 1) {
    hu[0] = fftw_malloc(sizeof(fftw_complex)*(in->n));
    hk[0] = fftw_malloc(sizeof(fftw_complex)*(in->n));
    hk[1] = fftw_malloc(sizeof(fftw_complex)*(in->n));

    ddk[0] = fftw_malloc(sizeof(fftw_complex)*(in->n));
    ddk[1] = fftw_malloc(sizeof(fftw_complex)*(in->n));

    Dyk = fftw_malloc(sizeof(fftw_complex)*(in->n));
    Rk = fftw_malloc(sizeof(fftw_complex)*(in->n));
    Dk = fftw_malloc(sizeof(fftw_complex)*(in->n));
    kD = fftw_malloc(sizeof(fftw_complex)*(in->n));
    D  = fftw_malloc(sizeof(fftw_complex)*(in->n));
    MiRk = fftw_malloc(sizeof(fftw_complex)*(in->n));
  }
  u  = fftw_malloc(sizeof(double)*(in->n));
  k  = fftw_malloc(sizeof(double)*(in->n));
  mk = fftw_malloc(sizeof(double)*(in->n));
  q  = (in->g)/pow(in->c,2);
  for (int j = 0; j < in->n; j ++) {
    u[j] = 2*pi*j/(in->n);
    k[j] = j;	
    if (j > (in->n)/2) k[j] = j - (in->n);
    mk[j] = cabs(1.*k[j]) - q;
    wrk->Y[j] = -0.4*cos(u[j]); 
  }
  fh = fopen(in->resname,"r");
  int jj = 0;
  if (fh != NULL) {
    char v1[80],v2[80],v3[80],v4[80], *dummy;
    while (fgets(line, 480, fh)!=NULL) {
      if (jj == 0) {
        dummy = fgets(line, 160, fh);
        dummy = fgets(line, 160, fh);
      }
      sscanf(line, "%s\t%s\t%s\t%s\n",  v1, v2, v3, v4);
      wrk->Y[jj] = atof(v4);
      jj++;
    }
    printf("Lines read = %d of %d expected\n", jj, N);
    if ( jj != N) err_msg("I/O error\n");
    fclose(fh);
  }

  //if (fftw_import_wisdom_from_filename("c2c.fftw") == 0) {
    //printf("Creating plans\n");
    //fftw_set_timelimit(60.);
    plan_f1 = fftw_plan_dft_1d(in->n, a, a, FFTW_FORWARD, fmode);
    plan_b1 = fftw_plan_dft_1d(in->n, a, a, FFTW_BACKWARD, fmode);
    //fftw_export_wisdom_to_filename("c2c.fftw");
    //err_msg("Plans created. Rerun the computation");
  //} else {
    //plan_f1 = fftw_plan_dft_1d(in->n, a, a, FFTW_FORWARD, fmode);
    //plan_b1 = fftw_plan_dft_1d(in->n, a, a, FFTW_BACKWARD, fmode);
  //}

  overN = 1./N; 
  in->w_ptr = wrk; dtau = in->dtau;
  gama = 1. + 1./dtau; tol = 1.e-15*N;
  fft(wrk->Y,wrk->Yk);
  for(int j = 0; j < N; j++) wrk->Yk[j] = (wrk->Yk[j])*overN*cexp((in->alfa)*cabs(1.*k[j])); 
  if (jj == 0) {
    wrk->Yk[0] = 0;
    for (int j = 1; j < N; j++) wrk->Yk[0] += -cabs(k[j])*conj(wrk->Yk[j])*(wrk->Yk[j]); 
  }
}

void khat(fftw_complex *in, fftw_complex *out) {
  memcpy(a, in, sizeof(fftw_complex)*N);
  fftw_execute(plan_f1);
  for (int j = 0; j < N ; j++) a[j] = cabs(1.*k[j])*overN*a[j];
  fftw_execute(plan_b1);
  memcpy(out, a, sizeof(fftw_complex)*N);  
}

void fft(fftw_complex *in, fftw_complex *out) {
  memcpy(a, in, N*sizeof(fftw_complex));
  fftw_execute(plan_f1);
  memcpy(out, a, N*sizeof(fftw_complex));
}

void ifft(fftw_complex *in, fftw_complex *out) {
  memcpy(a, in, N*sizeof(fftw_complex));
  fftw_execute(plan_b1);
  memcpy(out, a, N*sizeof(fftw_complex));
}

void prepare_rhs(work_ptr wrk) {
  ifft(wrk->Yk,wrk->Y); //ml0 = 0;
  for (int j = 0; j < N; j++) {
    hk[0][j] = (wrk->Yk[j])*cabs(k[j]);
    hk[1][j] = (wrk->Y[j])*(wrk->Y[j]);
    //ml0 += conj(wrk->Yk[j])*(hk[0][j]);
  }
  fft(hk[1],hk[1]);  
  ifft(hk[0],hu[0]);
  for (int j = 0; j < N; j++) {
    hk[1][j] = cabs(k[j])*hk[1][j]*overN;   
    hk[0][j] = hu[0][j]*(wrk->Y[j])*overN;
  }
  fft(hk[0],hk[0]); 
  for (int j = 0; j < N; j++) {
    hk[0][j] = -mk[j]*(wrk->Yk[j]) + q*(0.5*hk[1][j] + hk[0][j]); 
  }
}

fftw_complex inner_product(fftw_complex *in1, fftw_complex *in2) {
  inner = 0.;
  for (int j = 0; j < N; j++) {
    inner += conj(in1[j])*in2[j];
  }
  return inner;
}

void prepare_linearization(work_ptr wrk) {
  ifft(Dk, D);
  for (int j = 0; j < N; j++) {
    ddk[0][j] = (wrk->Y[j])*D[j];   
    ddk[1][j] = cabs(k[j])*Dk[j];    
  }
  fft(ddk[0], ddk[0]);
  ifft(ddk[1], kD);
  for (int j = 0; j < N; j++) {
    ddk[0][j] = cabs(k[j])*ddk[0][j]*overN;
    ddk[1][j] = (wrk->Y[j])*kD[j] + hu[0][j]*D[j];
  }
  fft(ddk[1], ddk[1]);
  for (int j = 0; j < N; j++) {
    ddk[1][j] = ddk[1][j]*overN;
    ddk[0][j] = -mk[j]*Dk[j] + q*(ddk[0][j] + ddk[1][j]);
  }
}

void iterate_cg(work_ptr wrk) {
  int counter, counter2;
  double cg_error = 1.;
  double err = 1.;
  double eps = 0.01;

  if ( fix_flag == 1 ) {
    for (int j = 0; j < N; j++) wrk->Yk[j]= (-1.i*signum(k[j]) + 1.i)*(wrk->Yk[j]);
    ifft(wrk->Yk,wrk->Y);
    write_out(wrk->Y,"stokes_phys.dat");
    write_out(wrk->Y,"restart.dat");
    write_out(wrk->Yk,"stokes_spec.dat");
    err_msg("Restored x-component.");
  }

  counter2 = 1; NV = 1.; counter = 0;
  //err_msg("Completed");
  while (NV > tol) {
    prepare_rhs(wrk);
    for (int j = 0; j < N; j++) {
      Dyk[j] = 0;
      Rk[j] = -1.*hk[0][j];
      Dk[j] = Rk[j]/mk[j];
    }
    NV = sqrt(creal(inner_product(hk[0],hk[0])) ); 
    counter = 0; cg_error = sqrt(creal(inner_product(Rk,Rk)));
    while (cg_error > eps*NV ) {       // start CG cycle 
      prepare_linearization(wrk);
      for (int j = 0; j < N; j++) MiRk[j] = Rk[j]/mk[j];
      num = inner_product(Rk,MiRk);
      den = inner_product(Dk,ddk[0]);
      aa = creal(num/den);
      for (int j = 0; j < N; j++) {
        Dyk[j] = Dyk[j] + aa*Dk[j];
        hk[1][j] = Rk[j];                       
        Rk[j]  = Rk[j] - aa*ddk[0][j];
        hk[0][j] = Rk[j]/mk[j];
      }
      num = inner_product(Rk, hk[0]);
      den = inner_product(hk[1], MiRk);
      bb = creal(num/den);
      for (int j = 0; j < N; j++) Dk[j] = hk[0][j] + bb*Dk[j];
      cg_error = sqrt(creal(inner_product(Rk,Rk)));
      counter++;
      if ((counter%400)==0) printf("Error = %e at %d step, a = %e  b = %e\n", cg_error*overN, counter, aa, bb);
      if (counter == 1000) break;

    }                                  // end of CG cycle
    if (counter == 1000) break;
    if ((counter2%1) == 0 ) {
        hmax = ((wrk->Y[N/2])-(wrk->Y[0]))/(2.*pi);
	printf("Nevyazka = %e Newton cycle %1d ", NV*overN, counter2);
	printf("Error = %e at %3d CG step, H/L = %e Speed = %.12e\n", cg_error*overN, counter, hmax, 1./sqrt(q));
    }
    for (int j = 0; j < N; j++) wrk->Yk[j] = wrk->Yk[j] + Dyk[j];
    counter2++;

  }  // end of outer cycle  

  if (counter == 1400) err_msg("CG diverged");

  for (int j = 0; j < N; j++) wrk->Yk[j]= (-1.i*signum(k[j]) + 1.i)*(wrk->Yk[j]);
  ifft(wrk->Yk,wrk->Y); 
  //for (int j = 0; j < N; j++) wrk->Y[j]= 1.i*(wrk->Y[j]);
  write_out(wrk->Y,"stokes_phys.dat");
  write_out(wrk->Y,"restart.dat");
  write_out(wrk->Yk,"stokes_spec.dat");
  FILE *fid = fopen("steep.data","a");
  fprintf(fid, "%.12e\t%.12e\n", hmax, 1./sqrt(q));
  fclose(fid);
  
}

/*void iterate(work_ptr wrk) {
  err = 1;
  int nn = 0;
  char st[80];

  write_out(wrk->Yk,"before.txt");
  while (err > tol) { 
    prepare_rhs(wrk); num = 0.; den = 0.;
    for (int j = 0; j < N; j++) {
      num += conj(wrk->Yk[j])*(hk[3][j]);
      den += mk[j]*pow(cabs(wrk->Yk[j]),2);
    }
    sp = 1.*num/den;
    err = 0.; l2 = 0; err2;
    for (int j = 0; j < N; j++) {
      rhs_yk = ((hk[3][j])/mk[j] - gama*sp*(wrk->Yk[j]) )*dtau;
      wrk->Yk[j] += rhs_yk;
      if (j > N/2) wrk->Yk[j] = conj(wrk->Yk[N-j]);
    }
    for (int j = 0; j < N; j++) {
      l2  += cpow(cabs(wrk->Yk[j]),2);
      err += cpow(cabs(rhs_yk),2);
      err2 += cpow(cabs(hk[3][j]),2);
    }
    err = sqrt(err/l2);
    err2 = sqrt(err2); 
    nn++;
    if ( (nn % 10)==0 ) {  
      printf("Iteration %5d. distance:  %e\tHeight:  %e\tSpeed: %.15e\tFluid Level: %e\n",nn, err,hmax/(2*pi), 1/sqrt(q),ml);
    }
  }
  printf("Iteration %5d. distance:  %e\tHeight:  %e\tSpeed: %.15e\tFluid Level: %e\n",nn, err,hmax/(2*pi), 1/sqrt(q),ml);
  ifft(wrk->Yk,wrk->Y);
  write_out(wrk->Y,"stokes_phys.dat");
  write_out(wrk->Y,"restart.dat");
  write_out(wrk->Yk,"stokes_spec.dat");

}*/





















